/**
 * File: BamFullmatchProcessor.java
 * Created by: mhaimel
 * Created on: Jun 17, 2011
 * CVS:  $Id: BamFullmatchProcessor.java 1.0 Jun 17, 2011 4:18:25 PM mhaimel Exp $
 */
package uk.ac.ebi.curtain.processor.load.impl;

import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import net.sf.samtools.SAMFileReader;
import net.sf.samtools.SAMFileReader.ValidationStringency;
import net.sf.samtools.SAMRecord;
import net.sf.samtools.SAMRecordIterator;

import org.apache.commons.lang.StringUtils;

import uk.ac.ebi.curtain.model.graph.Arc;
import uk.ac.ebi.curtain.model.graph.NodeEdge;
import uk.ac.ebi.curtain.model.graph.curtain.ContigInfo;
import uk.ac.ebi.curtain.model.graph.curtain.CurtainGraph;
import uk.ac.ebi.curtain.model.graph.curtain.ReadWrapper;
import uk.ac.ebi.curtain.model.graph.curtain.impl.FillCurtainGraph;
import uk.ac.ebi.curtain.processor.AbstractProcessor;
import uk.ac.ebi.curtain.service.IdentifierStore;
import uk.ac.ebi.curtain.storage.sequence.SequenceBox;
import uk.ac.ebi.curtain.storage.sequence.fast.FastSequenceBoxWriter;
import uk.ac.ebi.curtain.utils.CurtainUncheckedException;
import uk.ac.ebi.curtain.utils.data.DataEntry;
import uk.ac.ebi.curtain.utils.data.FileType;
import uk.ac.ebi.curtain.utils.data.ReadType;
import uk.ac.ebi.curtain.utils.file.FileInfo;
import uk.ac.ebi.curtain.utils.io.impl.FileIO;
import uk.ac.ebi.curtain.utils.match.MatchDetails;
import uk.ac.ebi.curtain.utils.match.sam.SamMatchParser.SamMatchDetails;
import uk.ac.ebi.velvet.model.Strand;

/**
 * @author mhaimel
 *
 */
public class BamFullmatchProcessor extends AbstractProcessor<DataEntry<File>, String>{

	private final AtomicReference<FastSequenceBoxWriter> sbWriter = new AtomicReference<FastSequenceBoxWriter>();
	private final AtomicReference<FastSequenceBoxWriter> sbArcWriter = new AtomicReference<FastSequenceBoxWriter>();
	private final AtomicInteger mappingQualityCutoff = new AtomicInteger(0); // no filtering

	@Override
	protected String buildReturnObject() {
		return "Done";
	}
	
	public int getMappingQualityCutoff() {
		return mappingQualityCutoff.get();
	}
	
	public void setMappingQualityCutoff(int cOff){
		this.mappingQualityCutoff.set(cOff);
	}

	@Override
	public void process() {
		DataEntry<File> dataEntries = getInput();
		if(!dataEntries.getMetaData().getFileType().isIn(FileType.bam)){
			throw new CurtainUncheckedException(
					"File not supported for loading: " + getInput()+ "!! Supported are: " + Arrays.toString(FileType.getMapTypes()));
		}
		loadFile(dataEntries);
	}

	private void loadFile(DataEntry<File> dataEntries) {
		File bamFile = dataEntries.getEntries()[0];
		if(dataEntries.getEntries().length > 1){
			throw new CurtainUncheckedException("Merge files per category!!!" + Arrays.toString(dataEntries.getEntries()));
		}
		getLog().debug("Load Bam file " + bamFile + " ...");
		CurtainGraph graph = getContext().getGraph();

		Integer category = dataEntries.getMetaData().getCategory();
		/************* insert size filter ****************/
		Integer ins = getContext().getArguments().getCategory2InsertLength().get(category);
		boolean paired = dataEntries.getMetaData().getReadType().isIn(ReadType.getPaired());
		
		FillCurtainGraph filler = new FillCurtainGraph(category, graph, ins);
		filler.setDistanceInfo(true);

		/************* file handing ****************/
		FileInfo info = new FileInfo(bamFile, dataEntries.getMetaData().getFileType(), dataEntries.getMetaData().getReadType());
		SAMFileReader bam = new SAMFileReader(bamFile);
		bam.setValidationStringency(ValidationStringency.SILENT);
//		SAMFileReader bamMate = new SAMFileReader(bamFile);
//		bamMate.setValidationStringency(ValidationStringency.SILENT);

		/************* contig names ****************/
		IdentifierStore<String, Integer> ctgIdx = getContext().getContigIndex();
		
		/************* statistics ****************/
		Long outOfInsert = 0l;
		Long fullStored = 0l;
		Long hangingStored = 0l;
		Long[] invalidCnt = new Long[]{0l,0l, 0l};
		Long[] badQualityCnt = new Long[]{0l,0l};
		Map<String, SAMRecord> buffer = new HashMap<String, SAMRecord>();
		try {
			SAMRecordIterator iter = bam.iterator();
				try{
					while(iter.hasNext()){
						SAMRecord next = iter.next();

						/************* Singe end data ****************/
						if(!paired){ // Single-end data
							if(isValidQuality(next)) {
								if(null != next.getReferenceName()){
									int ctgId = ctgIdx.getIdentifier(next.getReferenceName());
									writeFullEntry(ctgId, next);
								}
							}
							continue;
						}

						/************* Paired-end data ****************/
						
						if(next.getProperPairFlag()){ 
//							if(!isValidInsert(min,max,next)){// TODO Not at the moment - see method for more details
//								if(next.getFirstOfPairFlag()){
//									// remove down to out of insert size
//									++outOfInsert; // only count once
//								}
//								continue;
//							}
							// Used buffer
							SAMRecord[] pair = orderBufferedRecords(next, buffer);
							if(pair == null) continue;
							if(!isValidQuality(pair)) {
								++badQualityCnt[0];
								continue; /** remove down to bad quality **/
							}
							/************* Same contig ****************/
							if(!pair[1].getReadUnmappedFlag() && pair[1].getReferenceName().equals(pair[0].getReferenceName())){
								++fullStored;
								Integer ctgId = ctgIdx.getIdentifier(next.getReferenceName());
								writeFullEntry(ctgId, pair[0], pair[1]);	
								
								ContigInfo nodeValue = graph.getNode(ctgId.longValue()).getValue();
								nodeValue.addSequence(category,(long) (pair[0].getReadLength()+pair[1].getReadLength()));// 2x51
								nodeValue.incrementFull(category);
							} else {
								// Should not happen - invalid entry
								++invalidCnt[0];
							}
						} else {
							SAMRecord[] pair = orderBufferedRecords(next, buffer);
							if(pair == null) continue;
							
							String a = pair[0].getReferenceName();
							String b = pair[1].getReferenceName();
							if(bothUnmapped(pair)){
								/************* both unmapped ****************/
								++invalidCnt[2];
								// both unmapped Ignore for the moment, but TODO for later
							} else if(!StringUtils.equals(a, b) && !next.getMateUnmappedFlag() && !next.getReadUnmappedFlag()){
								/************* Different contig, both mapped ****************/
								if(isMultiBlockMap(pair)){
									continue; // Ignore multi block mapping for short reads on the edges TODO change for long reads
								}
								Integer aId = getContext().getContigIndex().getIdentifier(a);
								Integer bId = getContext().getContigIndex().getIdentifier(b);
								SamMatchDetails aDetail = new SamMatchDetails(info, pair[0]);
								SamMatchDetails bDetail = new SamMatchDetails(info, pair[1]);
								Arc<ContigInfo, ReadWrapper> arc = filler.doMap(aId.longValue(), aDetail.getAsAlternativeMatch(), bId.longValue(), bDetail.getAsAlternativeMatch());
								writeDifferentContigEntry(arc.getValue().getId(), pair[0],pair[1]);	
							} else if((next.getMateUnmappedFlag() || next.getReadUnmappedFlag())){
								/************* One mapped, other unmapped ****************/
								// Used buffer
//								SAMRecord[] pair = orderBufferedRecords(next, buffer);
//								if(pair == null) continue;
								next = null;
								next = pair[0].getReadUnmappedFlag()?pair[1]:pair[0];
								if(null == next || !isValidQuality(next)) {
									++badQualityCnt[1];
									continue; /** remove down to bad quality **/
								}
								MatchDetails details = new SamMatchDetails(info, pair[0]);
								
								++hangingStored;
								Integer ctgId = ctgIdx.getIdentifier(next.getReferenceName());
								writeFullEntry(ctgId, pair[0], pair[1]);
								ContigInfo nodeValue = graph.getNode(ctgId.longValue()).getValue();
								nodeValue.addSequence(category,(long) (pair[0].getReadLength()));// 1x51 - other not mapped
								nodeValue.incrementPartial(category, NodeEdge.fromBoolean(Strand.FWD.equals(details.getAlignOrientation())));
							} else {
								// Should not happen - invalid entry
								++invalidCnt[1];
							}
						} 
					}
				}finally {
					iter.close();
				}
//			}	
		} catch (IOException e) {
			throw new CurtainUncheckedException(e);
		}finally{
			FileIO.closeQuietly(bam);
		}
		getLog().debug("Finished loading Bam file " + bamFile + "!");
		getLog().debug("Invalid count for " + bamFile + ":" + Arrays.toString(invalidCnt));
		getLog().info("Stored hits: "+fullStored+"; Hanging hit stored: "+ hangingStored+"; Filtered by InsertSize: " + outOfInsert);
	}
	
	private boolean bothUnmapped(SAMRecord[] pair) {
		for(SAMRecord sr : pair){
			Integer idx = sr.getReferenceIndex();
			if(idx >= 0){
				return false;
			}
		}
		return true;
	}

	private void writeDifferentContigEntry(Long id,SAMRecord a,SAMRecord b) throws IOException {
		this.sbArcWriter.get().write(
				id,new SequenceBox(a.getReadName(), a.getReadString(),b.getReadString()));
	}

	private boolean isMultiBlockMap(SAMRecord ... reads) {
		for(SAMRecord r : reads){
			if(r.getAlignmentBlocks().size() > 1 ){
				return true; // only the first one - bad short read mapping
			}
		}
		return false;
	}

	private boolean isValidQuality(SAMRecord ... read) {
		for(SAMRecord sr : read){
			if(sr.getMappingQuality() < getMappingQualityCutoff()){
				return false;
			}
		}
		return true;
	}

	private boolean isValidInsert(Integer min, Integer max, SAMRecord next) {
		// TODO Introduce later for Mate-pair reads
		int infInsert = Math.abs(next.getInferredInsertSize()); 
		if(infInsert >= min && infInsert <= max){
			return true;
		}
		return false;
	}

	private SAMRecord[] orderBufferedRecords(SAMRecord next, Map<String, SAMRecord> buffer) {
		SAMRecord other = buffer.remove(next.getReadName());
		if(null == other){
			buffer.put(next.getReadName(), next);
			return null;
		}
		if(!next.getFirstOfPairFlag()){
			if(!other.getFirstOfPairFlag()){
				throw new IllegalStateException("Both entries first of pair NOT possible!!" + next);
			}
			return new SAMRecord[]{other,next};
		}	
		return new SAMRecord[]{next,other};
	}

//	private void initContigNames(String[] names) {
//		for(String s : getContext().getContigIndex().getKeys()){
//			Integer id = getContext().getContigIndex().getIdentifier(s);
//			names[id] = s;
//		}
//	}

	private void writeFullEntry(int i, SAMRecord next, SAMRecord other) throws IOException {
		this.sbWriter.get().write(
			(long)i,new SequenceBox(next.getReadName(), next.getReadString(),other.getReadString()));
	}
	
	private void writeFullEntry(int i, SAMRecord next) throws IOException {
		this.sbWriter.get().write(
				(long)i, new SequenceBox(next.getReadName(), next.getReadString()));
	}

	public void setSequenceBoxWriter(FastSequenceBoxWriter sbWriter) {
		this.sbWriter.set(sbWriter);
	}
	public FastSequenceBoxWriter getSbWriter() {
		return sbWriter.get();
	}
	
	public void setSequenceBoxArcWriter(FastSequenceBoxWriter sbWriter) {
		this.sbArcWriter.set(sbWriter);
	}
	public FastSequenceBoxWriter getSbArcWriter() {
		return sbArcWriter.get();
	}
}









